﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using CommModule.Security;
using CommModule.Messages;
using System.Security.Cryptography;

namespace CommModule.Connection
{
    public class TCPSecureSocket
    {
        private enum ProtocolPhases
        {
            WaitingSessionKeyMessage, WaitingServerChallengeAnswer
        };

        private readonly bool _bypass = false;
        private bool _secureChannelEstablished;
        private bool _establishingSecureChannel;

        private Object _lastChannelMessage;
        private String _lastReceivedMessageTS = "0";

        private byte[] _sessionKey;

        private TCPSocket _tcpSocket;

        public TCPSecureSocket()
        {
            _tcpSocket = new TCPSocket();

            _establishingSecureChannel = true;
            _secureChannelEstablished = EstablishSecureChannel();
        }

        public Object Receive()
        {
            Object obj = _tcpSocket.Receive();

            if (obj == null)
                return null;

            if (_bypass || _establishingSecureChannel)
                return obj;
            else if(!obj.GetType().Name.Equals("CipheredMessage"))
                return null;

            String cipheredString = ((CipheredMessage)obj).CipheredMessageString;
            String secureMessageString = SimmEncryption.Decipher(cipheredString, _sessionKey);
            SecureMessage secureMessage = (SecureMessage)ObjectSerialization.DeserializeObjectFromXml(Encoding.UTF8.GetBytes(secureMessageString));

            RSACryptoServiceProvider certcopPubKeyDecryptor = CertificateUtils.GetCertcopPublicKey();
            bool validMessage = secureMessage.ValidateMessage(_lastReceivedMessageTS, certcopPubKeyDecryptor);

            if (validMessage)
            {
                _lastReceivedMessageTS = secureMessage.GetTimestamp();
                return secureMessage.GetObject();
            }
            else
                return null;
        }

        public void Send(Object obj)
        {
            if (_bypass || _establishingSecureChannel)
                _tcpSocket.Send(obj);
            else
            {
                RSACryptoServiceProvider serverPrivateKeyEncryptor = CertificateUtils.GetServerPrivateKey();
                SecureMessage secureMessage = new SecureMessage(obj, serverPrivateKeyEncryptor);
                String secureMessageString = ObjectSerialization.SerializeObjectAsString(secureMessage);
                String cipheredString = SimmEncryption.Cipher(secureMessageString, _sessionKey);

                CipheredMessage cipheredMessage = new CipheredMessage(cipheredString);
                _tcpSocket.Send(cipheredMessage);
            }
        }

        private bool EstablishSecureChannel()
        {
            int nounce = SecureRandomNumber.GenerateNounce();
            CertChallengeMessage challengeMessage = new CertChallengeMessage(nounce);
            _lastChannelMessage = challengeMessage;
            Send(challengeMessage);

            CertSKExchangeMessage skMessage = (CertSKExchangeMessage)ReceiveProtocolMessage(ProtocolPhases.WaitingSessionKeyMessage);
            _sessionKey = skMessage.GetSessionKey();
            int serverChallenge = skMessage.GetClientChallenge();

            int clientChallenge = SecureRandomNumber.GenerateNounce();
            SecChannelChallengeMessage challengePair = new SecChannelChallengeMessage(clientChallenge, serverChallenge, _sessionKey);
            _lastChannelMessage = challengePair;
            Send(challengePair);

            SecChannelChallengeMessage serverAnswer = (SecChannelChallengeMessage)ReceiveProtocolMessage(ProtocolPhases.WaitingServerChallengeAnswer);
            int serverChallengeAnswer = serverAnswer.GetChallenge(_sessionKey);

            if (serverChallengeAnswer == clientChallenge)
            {
                _establishingSecureChannel = false;
                return true;
            }

            return false;
        }

        public bool IsSecureChannelEstablished()
        {
            return _secureChannelEstablished;
        }

        private Object ReceiveProtocolMessage(ProtocolPhases phase)
        {
            Object obj = null;
            while (obj == null)
                obj = Receive();

            if (obj == null)
                return null;
            else if (obj.GetType().Name.Equals("CertSKExchangeMessage") && phase == ProtocolPhases.WaitingSessionKeyMessage)
                return obj;
            else if (obj.GetType().Name.Equals("SecChannelChallengeMessage") && phase == ProtocolPhases.WaitingServerChallengeAnswer)
                return obj;
            else
            {
                Send(new UnexpectedMessage());
                return ReceiveProtocolMessage(phase);
            }
        }
    }
}
